def plot_wigner2(rho, fig=None, ax=None, figsize=(6, 6),
cmap=None, alpha_max=7.5, colorbar=False,
method='clenshaw', projection='2d'):
"""
Plot the the Wigner function for a density matrix (or ket) that describes
an oscillator mode.
Parameters
----------
rho : :class:`qutip.qobj.Qobj`
The density matrix (or ket) of the state to visualize.
fig : a matplotlib Figure instance
The Figure canvas in which the plot will be drawn.
ax : a matplotlib axes instance
The axes context in which the plot will be drawn.
figsize : (width, height)
The size of the matplotlib figure (in inches) if it is to be created
(that is, if no 'fig' and 'ax' arguments are passed).
cmap : a matplotlib cmap instance
The colormap.
alpha_max : float
The span of the x and y coordinates (both [-alpha_max, alpha_max]).
colorbar : bool
Whether (True) or not (False) a colorbar should be attached to the
Wigner function graph.
method : string {'clenshaw', 'iterative', 'laguerre', 'fft'}
The method used for calculating the wigner function. See the
documentation for qutip.wigner for details.
projection: string {'2d', '3d'}
Specify whether the Wigner function is to be plotted as a
contour graph ('2d') or surface plot ('3d').
Returns
-------
fig, ax : tuple
A tuple of the matplotlib figure and axes instances used to produce
the figure.
"""
if not fig and not ax:
if projection == '2d':
pass
# fig, ax = plt.subplots(1, 1, figsize=figsize)
elif projection == '3d':
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1, projection='3d')
else:
raise ValueError('Unexpected value of projection keyword argument')
if isket(rho):
rho = ket2dm(rho)
xvec = np.linspace(-alpha_max, alpha_max, 200)
W0 = wigner(rho, xvec, xvec, method=method)
W, yvec = W0 if type(W0) is tuple else (W0, xvec)
wlim = abs(W).max()
if cmap is None:
cmap = cm.get_cmap('RdBu')
if projection == '2d':
# cf = ax.contourf(xvec, yvec, W, 100,
# norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
# cf = go.Figure(xvec, yvec, W)
cf = go.Figure(data=
go.Heatmap(
z=W, connectgaps=True))
cf.show()
elif projection == '3d':
X, Y = np.meshgrid(xvec, xvec)
cf = ax.plot_surface(X, Y, W0, rstride=5, cstride=5, linewidth=0.5,
norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
else:
raise ValueError('Unexpected value of projection keyword argument.')
# if xvec is not yvec:
# ax.set_ylim(xvec.min(), xvec.max())
# ax.set_xlabel(r'$\rm{Re}(\alpha)$', fontsize=12)
# ax.set_ylabel(r'$\rm{Im}(\alpha)$', fontsize=12)
# if colorbar:
# fig.colorbar(cf, ax=ax)
#
# ax.set_title("Wigner function", fontsize=12)
# return fig, ax